# coding: utf-8
import numpy as np
import matplotlib.pyplot as plt
from random import shuffle, uniform
from copy import deepcopy

np.set_printoptions(suppress=True, linewidth=999, threshold=np.inf)
plt.rc('font', family='Times New Roman', weight='normal', size=12)


class Solution:
    def __init__(self):
        self.tours = None
        self.costs = None
        self.total = None


class ABC:
    def __init__(self, data):
        self.V = data.cities  # city list
        self.n = len(self.V)  # the number of cities
        self.Distance = data.distance_matrix  # distance matrix

        self.Depot = 1
        self.m = 2
        self.Pr = 0.8

        self.optima = None  # the best ant ever in history
        self.evolve_list = []  # optimal cost in evolution progress

    def init_solution(self, solution):
        solution.tours = [[self.Depot, self.Depot] for _ in range(self.m)]
        solution.costs = [0 for _ in range(self.m)]
        cities = deepcopy(self.V)
        cities.remove(self.Depot)
        shuffle(cities)
        for city in cities:
            solution = self.optimal_insert(solution, city)
        solution = self.calc_costs(solution)
        return solution

    def calc_costs(self, solution):
        for r, tour in enumerate(solution.tours):
            solution.costs[r] = 0
            for i in range(1, len(tour)):
                solution.costs[r] += self.Distance[tour[i - 1]][tour[i]]
        solution.total = sum(solution.costs)
        return solution

    def optimal_insert(self, solution, city):
        assert solution.tours is not None
        opt_r, opt_i, opt_cost = -1, -1, 1e99
        r = np.argmin(np.array([len(tour) for tour in solution.tours]))

        # for r in range(len(solution.tours)):
        for i in range(1, len(solution.tours[r])):
            p, q = solution.tours[r][i - 1], solution.tours[r][i]
            cost = self.Distance[p][city] + self.Distance[city][q] - self.Distance[p][q]
            if opt_cost > cost:
                opt_r, opt_i, opt_cost = r, i, cost
        solution.tours[opt_r].insert(opt_i, city)
        solution.costs[opt_r] += opt_cost
        solution.total = sum(solution.costs)
        return solution

    def optimize(self, iteration=1000):
        self.optima = Solution()
        self.optima.total = np.inf
        solution = self.init_solution(Solution())
        for it in range(1, iteration + 1):
            cities = []
            for tour in solution.tours:
                for i in range(len(tour) - 2, 0, -1):
                    if uniform(0, 1) < self.Pr:
                        cities.append(tour[i])
                        tour.remove(tour[i])
            shuffle(cities)
            for city in cities:
                solution = self.optimal_insert(solution, city)
            solution = self.calc_costs(solution)
            if self.optima.total > solution.total:
                self.optima = deepcopy(solution)
            self.evolve_list.append(self.optima.total)
            plt.clf()
            plt.subplot(1, 2, 1)
            plt.plot(range(len(self.evolve_list)), self.evolve_list)
            plt.subplot(1, 2, 2)
            data.draw(plt, routes=self.optima.tours)
            plt.title(f'{it}: {self.optima.total:.2f}')
            plt.pause(0.001)


from optimizer.mtsp.TSP import TSP
data = TSP(path=f'../../data/TSPLIB/att532.tsp.txt')
optimizer = ABC(data)
optimizer.optimize()


